/*
 * Wazuh storeRemediationsModel
 * Copyright (C) 2015, Wazuh Inc.
 * October 05, 2023.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#include "updateCVERemediations_test.hpp"
#include "flatbuffers/idl.h"
#include "flatbuffers/util.h"
#include <memory>

constexpr auto COMMON_DATABASE_DIR {"queue/vd"}; //<<Used for all databases
const std::string FLATBUFFER_SCHEMA {FLATBUFFER_SCHEMAS_DIR "/cve5.fbs"};
const char* FB_INCLUDE_DIRECTORIES[] = {FLATBUFFER_SCHEMAS_DIR, nullptr};

const auto JSON_CVE5_VALID_ONE_BLOCK = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-1337-1234",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "anyOf":["KBT-800","KBT-1000","KBT-3000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

const auto JSON_CVE5_VALID_MULTIPLE_BLOCKS = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-1337-1234",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "anyOf":["KBT-800","KBT-1000","KBT-3000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-4000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-5000","KBT-6000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-7000","KBT-8000","KBT-9000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

const auto JSON_CVE5_VALID_MULTIPLE_BLOCKS_SOME_WITH_NO_ANYOF = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-1337-1234",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "anyOf":["KBT-800","KBT-1000","KBT-3000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-4000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-7000","KBT-8000","KBT-9000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

const auto JSON_CVE5_EMPTY_UPDATES = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-1337-1234",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

const auto JSON_CVE5_NO_REMEDIATIONS = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-1337-1234",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
            }
        }
    }
)";

void UpdateCVERemediationsTest::SetUp()
{
    std::filesystem::create_directories(COMMON_DATABASE_DIR);
};

void UpdateCVERemediationsTest::TearDown()
{
    std::filesystem::remove_all(COMMON_DATABASE_DIR);
};

TEST_F(UpdateCVERemediationsTest, UpdatesWindowsRemediationOneBlock)
{
    // Define schema variable and parse JSON object.
    std::string schemaStr;

    // Load file with schema.
    bool valid = flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);
    EXPECT_TRUE(valid);

    // Parse schema.
    flatbuffers::Parser parser;
    valid = parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES) && parser.Parse(JSON_CVE5_VALID_ONE_BLOCK);
    EXPECT_TRUE(valid);

    // Create a test Entry object with Windows remediations
    auto jbuf = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier jverifier(jbuf, parser.builder_.GetSize());
    EXPECT_TRUE(cve_v5::VerifyEntryBuffer(jverifier));
    auto entry = cve_v5::GetEntry(jbuf);

    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateCVERemediations::storeVulnerabilityRemediation(entry, rocksDBWrapper.get()));

    // Assert result.
    std::string cveId {"CVE-1337-1234"};
    FlatbufferDataPair<RemediationInfo> dtoVulnRemediation;

    // Asserts
    EXPECT_NO_THROW(rocksDBWrapper->get(cveId, dtoVulnRemediation.slice, REMEDIATIONS_COLUMN));
    dtoVulnRemediation.data = const_cast<NSVulnerabilityScanner::RemediationInfo*>(
        NSVulnerabilityScanner::GetRemediationInfo(dtoVulnRemediation.slice.data()));

    ASSERT_EQ(dtoVulnRemediation.data->updates()->size(), 3);
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(0)->str().c_str(), "KBT-800");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(1)->str().c_str(), "KBT-1000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(2)->str().c_str(), "KBT-3000");
}

TEST_F(UpdateCVERemediationsTest, UpdatesWindowsRemediationMultipleBlocks)
{
    // Define schema variable and parse JSON object.
    std::string schemaStr;

    // Load file with schema.
    bool valid = flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);
    EXPECT_TRUE(valid);

    // Parse schema.
    flatbuffers::Parser parser;
    valid = parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES) && parser.Parse(JSON_CVE5_VALID_MULTIPLE_BLOCKS);
    EXPECT_TRUE(valid);

    // Create a test Entry object with Windows remediations
    auto jbuf = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier jverifier(jbuf, parser.builder_.GetSize());
    EXPECT_TRUE(cve_v5::VerifyEntryBuffer(jverifier));
    auto entry = cve_v5::GetEntry(jbuf);

    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateCVERemediations::storeVulnerabilityRemediation(entry, rocksDBWrapper.get()));

    // Assert result.
    std::string cveId {"CVE-1337-1234"};
    FlatbufferDataPair<RemediationInfo> dtoVulnRemediation;

    // Asserts
    EXPECT_NO_THROW(rocksDBWrapper->get(cveId, dtoVulnRemediation.slice, REMEDIATIONS_COLUMN));
    dtoVulnRemediation.data = const_cast<NSVulnerabilityScanner::RemediationInfo*>(
        NSVulnerabilityScanner::GetRemediationInfo(dtoVulnRemediation.slice.data()));

    ASSERT_EQ(dtoVulnRemediation.data->updates()->size(), 9);
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(0)->str().c_str(), "KBT-800");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(1)->str().c_str(), "KBT-1000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(2)->str().c_str(), "KBT-3000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(3)->str().c_str(), "KBT-4000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(4)->str().c_str(), "KBT-5000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(5)->str().c_str(), "KBT-6000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(6)->str().c_str(), "KBT-7000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(7)->str().c_str(), "KBT-8000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(8)->str().c_str(), "KBT-9000");
}

TEST_F(UpdateCVERemediationsTest, UpdatesWindowsRemediationMultipleBlocksSomeWithNoAnyOf)
{
    // Define schema variable and parse JSON object.
    std::string schemaStr;

    // Load file with schema.
    bool valid = flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);
    EXPECT_TRUE(valid);

    // Parse schema.
    flatbuffers::Parser parser;
    valid = parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES) &&
            parser.Parse(JSON_CVE5_VALID_MULTIPLE_BLOCKS_SOME_WITH_NO_ANYOF);
    EXPECT_TRUE(valid);

    // Create a test Entry object with Windows remediations
    auto jbuf = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier jverifier(jbuf, parser.builder_.GetSize());
    EXPECT_TRUE(cve_v5::VerifyEntryBuffer(jverifier));
    auto entry = cve_v5::GetEntry(jbuf);

    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateCVERemediations::storeVulnerabilityRemediation(entry, rocksDBWrapper.get()));

    // Assert result.
    std::string cveId {"CVE-1337-1234"};
    FlatbufferDataPair<RemediationInfo> dtoVulnRemediation;

    // Asserts
    EXPECT_NO_THROW(rocksDBWrapper->get(cveId, dtoVulnRemediation.slice, REMEDIATIONS_COLUMN));
    dtoVulnRemediation.data = const_cast<NSVulnerabilityScanner::RemediationInfo*>(
        NSVulnerabilityScanner::GetRemediationInfo(dtoVulnRemediation.slice.data()));

    ASSERT_EQ(dtoVulnRemediation.data->updates()->size(), 7);
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(0)->str().c_str(), "KBT-800");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(1)->str().c_str(), "KBT-1000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(2)->str().c_str(), "KBT-3000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(3)->str().c_str(), "KBT-4000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(4)->str().c_str(), "KBT-7000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(5)->str().c_str(), "KBT-8000");
    EXPECT_STREQ(dtoVulnRemediation.data->updates()->Get(6)->str().c_str(), "KBT-9000");
}

TEST_F(UpdateCVERemediationsTest, SkipsEmptyUpdates)
{
    // Define schema variable and parse JSON object.
    std::string schemaStr;

    // Load file with schema.
    bool valid = flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);
    EXPECT_TRUE(valid);

    // Parse schema.
    flatbuffers::Parser parser;
    valid = parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES) && parser.Parse(JSON_CVE5_EMPTY_UPDATES);
    EXPECT_TRUE(valid);

    // Create a test Entry object with Windows remediations
    auto jbuf = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier jverifier(jbuf, parser.builder_.GetSize());
    EXPECT_TRUE(cve_v5::VerifyEntryBuffer(jverifier));
    auto entry = cve_v5::GetEntry(jbuf);

    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    if (!rocksDBWrapper->columnExists(REMEDIATIONS_COLUMN))
    {
        rocksDBWrapper->createColumn(REMEDIATIONS_COLUMN);
    }

    // Asserts there's no element with the key provided.
    EXPECT_NO_THROW(UpdateCVERemediations::storeVulnerabilityRemediation(entry, rocksDBWrapper.get()));
    std::string key {entry->cveMetadata()->cveId()->str()};
    FlatbufferDataPair<RemediationInfo> dtoVulnRemediation;
    EXPECT_EQ(false, rocksDBWrapper->get(key, dtoVulnRemediation.slice, REMEDIATIONS_COLUMN));
}

TEST_F(UpdateCVERemediationsTest, x_remediationsNotPresent)
{
    // Define schema variable and parse JSON object.
    std::string schemaStr;

    // Load file with schema.
    bool valid = flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);
    EXPECT_TRUE(valid);

    // Parse schema.
    flatbuffers::Parser parser;
    valid = parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES) && parser.Parse(JSON_CVE5_NO_REMEDIATIONS);
    EXPECT_TRUE(valid);

    // Create a test Entry object with Windows remediations
    auto jbuf = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier jverifier(jbuf, parser.builder_.GetSize());
    EXPECT_TRUE(cve_v5::VerifyEntryBuffer(jverifier));
    auto entry = cve_v5::GetEntry(jbuf);

    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    if (!rocksDBWrapper->columnExists(REMEDIATIONS_COLUMN))
    {
        rocksDBWrapper->createColumn(REMEDIATIONS_COLUMN);
    }

    // Asserts there's no element with the key provided.
    EXPECT_NO_THROW(UpdateCVERemediations::storeVulnerabilityRemediation(entry, rocksDBWrapper.get()));
    std::string key {entry->cveMetadata()->cveId()->str()};
    FlatbufferDataPair<RemediationInfo> dtoVulnRemediation;
    EXPECT_EQ(false, rocksDBWrapper->get(key, dtoVulnRemediation.slice, REMEDIATIONS_COLUMN));
}

TEST_F(UpdateCVERemediationsTest, DeleteRemediation)
{
    // Define schema variable and parse JSON object.
    std::string schemaStr;

    // Load file with schema.
    bool valid = flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);
    EXPECT_TRUE(valid);

    // Parse schema.
    flatbuffers::Parser parser;
    valid = parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES) && parser.Parse(JSON_CVE5_VALID_ONE_BLOCK);
    EXPECT_TRUE(valid);

    // Create a test Entry object with Windows remediations
    auto jbuf = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier jverifier(jbuf, parser.builder_.GetSize());
    EXPECT_TRUE(cve_v5::VerifyEntryBuffer(jverifier));
    auto entry = cve_v5::GetEntry(jbuf);

    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateCVERemediations::storeVulnerabilityRemediation(entry, rocksDBWrapper.get()));

    // Make sure the entry was inserted
    FlatbufferDataPair<RemediationInfo> dtoVulnRemediation;
    std::string cveId {"CVE-1337-1234"};
    EXPECT_TRUE(rocksDBWrapper->get(cveId, dtoVulnRemediation.slice, REMEDIATIONS_COLUMN));

    // Call the removeRemediation function
    EXPECT_NO_THROW(UpdateCVERemediations::removeRemediation(entry, rocksDBWrapper.get()));

    // Asserts there's no element with the key provided.
    EXPECT_FALSE(rocksDBWrapper->get(cveId, dtoVulnRemediation.slice, REMEDIATIONS_COLUMN));
}
